import tensorflow.compat.v1 as tf
tf.disable_eager_execution()

'''
N.B. Operations that mutate a variable, 
such as v.assign() and the parameter update operations in a tf.train.Optimizer must run on the same device as the variable. 
Incompatible device placement directives will be ignored when creating these operations.
'''

if __name__ == '__main__':
    with tf.device('/job:ps/task:0/cpu:0'):
        W1 = tf.Variable(tf.ones(()),name='W1')
        b1 = tf.Variable(tf.ones(()),name='b1')
    
    with tf.device('/job:ps/task:1/cpu:0'):
        W2 = tf.Variable(tf.ones(()),name='W2')
        b2 = tf.Variable(tf.ones(()),name='b2')
    
    with tf.device('/job:worker/task:0/gpu:0'):
        A1 = tf.add(tf.multiply(1.0,W1),b1)
    
    with tf.device('/job:ps/task:0/cpu:0'):
        At = A1-0.5
        
    with tf.device('/job:worker/task:1/gpu:0'):
        A1 = At*A1
        A2 = tf.add(tf.multiply(A1,W2),b2)
        cost = tf.reduce_sum(tf.pow(A2-1,2))
        # cost = tf.reduce_sum(tf.pow(A2-1.0,2))
        
    with tf.device('/job:ps/task:0/cpu:0'):
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.1)
        train_op = optimizer.minimize(cost)

    with tf.Session('grpc://192.168.2.43:2222',
                    config=tf.ConfigProto(
                        log_device_placement=True
                        # ,
                        # graph_options=tf.GraphOptions(
                        #     optimizer_options=tf.OptimizerOptions(
                        #         opt_level=tf.OptimizerOptions.L0)
                        #     )
                        #
                        )
                    ) as sess:
        sess.run(tf.global_variables_initializer())
        run_options = tf.RunOptions(
            trace_level=tf.RunOptions.FULL_TRACE
        )
        run_metadata = tf.RunMetadata()
        for i in range(10):
            # loss0, a0, b0, _ = sess.run([loss, a, b, train_op])
            # print(f"loss={loss0}, a={a0}, b={b0}")
            ret = sess.run([cost, W1, b1, W2, b2, train_op],
                           options=run_options,
                           run_metadata=run_metadata)
            print(f"ret={ret}")
        writer = tf.summary.FileWriter('board', sess.graph)
        writer.close()  
        with open('./op_placement.txt','w') as f:
            for device in run_metadata.step_stats.dev_stats:
                f.write(f"{device.device}:\t")
                for node in device.node_stats:
                    f.write(f"{node.node_name}\t")
                f.write("\n\n")